import cv2
import numpy as np
import math
import pytesseract
import re
import matplotlib.pyplot as plt
from segment_anything import SamPredictor, sam_model_registry
import pytesseract

pytesseract.pytesseract.tesseract_cmd = r'D:\Tesseract-OCR\tesseract.exe'

# 提取图像中比例尺的函数，使用OCR技术
def extract_scale_from_image(image):
    """
    处理图像并使用OCR提取比例尺。

    :param image_path: 需要处理的图像路径
    :return: 提取的比例尺，格式为字符串
    """
    # 读取图像
    # image = cv2.imread(image_path)

    # 获取图像的尺寸，并裁剪右下角区域
    height, width, _ = image.shape
    crop_x_start = int(width * 0.75)
    crop_y_start = int(height * 0.75)
    cropped_image = image[crop_y_start:height, crop_x_start:width]

    # 将图像转换为灰度图，应用高斯模糊，然后进行二值化处理
    gray = cv2.cvtColor(cropped_image, cv2.COLOR_BGR2GRAY)
    gray = cv2.GaussianBlur(gray, (5, 5), 0)
    _, thresh = cv2.threshold(gray, 127, 255, cv2.THRESH_BINARY)

    # 使用OCR从处理后的图像中提取文本
    custom_config = r'--oem 3 --psm 6'
    text = pytesseract.image_to_string(thresh, config=custom_config)

    # 清理和处理OCR提取的文本
    text = text.replace('um', 'μm')  # 将'um'替换为'μm'
    text = re.sub(r'[^\w\sμm]', '', text)  # 移除不必要的字符
    text = text.replace("\n", " ")  # 将换行符替换为空格

    # 使用正则表达式提取数字后跟有效单位的部分
    filtered_text = re.findall(r'(\d+\.?\d*)\W*(μm|mm|m)', text)

    if filtered_text:
        # 将结果连接成一个字符串，例如："1000μm"
        final_text = "".join(["".join(match) for match in filtered_text]).strip()
    else:
        final_text = "未识别到有效数据"  # 如果没有有效数据，返回默认消息

    return final_text

def extract_scale_bar(source_image):
    """提取图像中的比例尺线段像素数量"""
    if source_image is None:
        return None, None, None

    try:
        # 将图像转换为灰度
        gray = cv2.cvtColor(source_image, cv2.COLOR_BGR2GRAY)

        # 对灰度图像进行二值化处理
        _, thresh = cv2.threshold(gray, 1, 255, cv2.THRESH_BINARY)

        # 检测图像中的所有轮廓
        contours, _ = cv2.findContours(thresh, cv2.RETR_TREE, cv2.CHAIN_APPROX_SIMPLE)

        # 初始化比例尺轮廓
        scale_contour = None

        # 图像尺寸
        img_height, img_width = source_image.shape[:2]

        # 遍历所有轮廓，查找比例尺
        for contour in contours:
            # 计算轮廓的边界框
            x, y, w, h = cv2.boundingRect(contour)

            area = cv2.contourArea(contour)

            # 检查轮廓是否可能是比例尺
            if area > 150:  # 面积过滤
                # 检查轮廓是否在图像边缘或底部区域
                if (x < 100 or y < 100 or
                        x + w > img_width - 100 or
                        y + h > img_height - 100):

                    # 偏向于选择图像下半部分的轮廓
                    if y > img_height / 2:
                        scale_contour = contour
                        break

        if scale_contour is None:
            return None, None, None

        # 提取包含比例尺的区域
        x, y, w, h = cv2.boundingRect(scale_contour)
        roi_img = source_image[max(0, y - 30):min(img_height, y + h + 40),
                  max(0, x - 20):min(img_width, x + w + 20)]

        # 尝试从比例尺区域提取像素数量
        try:
            # 将ROI转换为灰度
            roi_gray = cv2.cvtColor(roi_img, cv2.COLOR_BGR2GRAY)

            # 对ROI区域进行OCR识别
            text = pytesseract.image_to_string(roi_img)

            # 使用正则表达式查找比例尺信息
            scale_match = re.search(r'(\d+)\s*(μm|um|mm|m)', text, re.IGNORECASE)

            if scale_match:
                scale_value = scale_match.group(1)
                scale_unit = scale_match.group(2).lower()

                # 测量比例尺长度（像素）
                # 使用轮廓宽度作为比例尺长度的估计值
                scale_pixel_length = w

                return scale_pixel_length, scale_value, scale_unit
            else:
                # 如果OCR无法识别文本，尝试直接测量比例尺线段
                # 对ROI进行二值化和边缘检测
                _, roi_thresh = cv2.threshold(roi_gray, 127, 255, cv2.THRESH_BINARY_INV)

                # 寻找最长的水平线
                lines = cv2.HoughLinesP(roi_thresh, 1, np.pi / 180, 50, minLineLength=30, maxLineGap=10)

                if lines is not None:
                    max_length = 0
                    for line in lines:
                        x1, y1, x2, y2 = line[0]
                        length = math.sqrt((x2 - x1) ** 2 + (y2 - y1) ** 2)
                        if length > max_length:
                            max_length = length

                    # 只返回像素长度
                    return int(max_length), None, None
                else:
                    # 使用轮廓宽度作为备选方案
                    return w, None, None

        except Exception as e:
            # 使用轮廓宽度作为备选方案
            return w, None, None

    except Exception as e:
        return None, None, None
